import torch


def check_gpu_activity(gpu_id=0):
    device = torch.device(f'cuda:{gpu_id}')
    print(f"GPU {gpu_id} 当前活动进程数: {torch.cuda.device_count()}")
    print(f"GPU {gpu_id} 显存占用: {torch.cuda.memory_allocated(device) / 1024 ** 3:.2f} GB")
    print(f"GPU {gpu_id} 最大显存: {torch.cuda.device(gpu_id).total_memory / 1024 ** 3:.2f}")


# 检查指定 GPU（例如第 3 块 GPU）
check_gpu_activity(gpu_id=3)